import numpy as np

# ================= 场景常量与初始条件 =================
g = 9.81  # m/s^2
R = 10.0  # 云团有效半径
v_sink = 3.0  # 云团中心起爆后下沉速度 (m/s)
valid_win = 20.0  # 每枚弹起爆后仅 20 s 有效
T = np.array([0, 200, 0])  # 真目标中心
Of = np.array([0, 0, 0])  # 假目标（原点）

# 导弹初始
M1_0 = np.array([20000, 0, 2000])
M2_0 = np.array([19000, 600, 2100])
M3_0 = np.array([18000, -600, 1900])

# 导弹速度向量（300 m/s 指向假目标）
u1 = (Of - M1_0) / np.linalg.norm(Of - M1_0)
u2 = (Of - M2_0) / np.linalg.norm(Of - M2_0)
u3 = (Of - M3_0) / np.linalg.norm(Of - M3_0)
v1 = 300 * u1
v2 = 300 * u2
v3 = 300 * u3

# 抵达假目标时间（近似上限）
t_hit1 = np.linalg.norm(Of - M1_0) / 300
t_hit2 = np.linalg.norm(Of - M2_0) / 300
t_hit3 = np.linalg.norm(Of - M3_0) / 300
t_max = max([t_hit1, t_hit2, t_hit3])  # 统一时间轴上限

# 5 架 UAV 初始（等高度直线匀速，航向/速度一经选定便固定）
FY = np.array([
    [17800, 0, 1800],  # FY1
    [12000, 1400, 1400],  # FY2
    [6000, -3000, 700],  # FY3
    [11000, 2000, 1800],  # FY4
    [13000, -2000, 1300]  # FY5
])
NUM_UAV = FY.shape[0]
v_min = 70
v_max = 140

# ================= 评估网格（可按需调快/调精）=================
angles = np.deg2rad(np.arange(-180, 180, 30))  # 航向角网格（相对 +x 逆时针）
speeds = np.array([70, 100, 130])  # 速度网格
t_grid = np.arange(2, t_max - 2, 2)  # 候选"起爆全局时刻" (s)
tau_set = np.array([2.4, 3.0, 3.6, 4.2])  # 候选引信延时 (s)
dt_eval = 0.03  # 评估时间步长 (s) ——可调：更小更精
MAX_TOTAL_SHOTS = 12  # 全局最多选用的弹数（<= 15，调大更猛）

# ================= 预计算三枚导弹的轨迹与视线量 =================
ts = np.arange(0, t_max - 1e-6, dt_eval)
N = len(ts)

# 轨迹
M1x = M1_0[0] + v1[0] * ts
M1y = M1_0[1] + v1[1] * ts
M1z = M1_0[2] + v1[2] * ts

M2x = M2_0[0] + v2[0] * ts
M2y = M2_0[1] + v2[1] * ts
M2z = M2_0[2] + v2[2] * ts

M3x = M3_0[0] + v3[0] * ts
M3y = M3_0[1] + v3[1] * ts
M3z = M3_0[2] + v3[2] * ts

# 视线向量 U = T - M(t) 及 UU=|U|^2
U1x = T[0] - M1x
U1y = T[1] - M1y
U1z = T[2] - M1z
UU1 = U1x ** 2 + U1y ** 2 + U1z ** 2

U2x = T[0] - M2x
U2y = T[1] - M2y
U2z = T[2] - M2z
UU2 = U2x ** 2 + U2y ** 2 + U2z ** 2

U3x = T[0] - M3x
U3y = T[1] - M3y
U3z = T[2] - M3z
UU3 = U3x ** 2 + U3y ** 2 + U3z ** 2

# ================= 贪心增量优化（联合三枚导弹） =================
# 当前联合"最小距离"数组（分别对三枚导弹），初始为 +Inf
dmin1 = np.full(N, np.inf)
dmin2 = np.full(N, np.inf)
dmin3 = np.full(N, np.inf)

# 每架 UAV 的状态
uav_fixed = np.zeros(NUM_UAV, dtype=bool)
uav_theta = np.zeros(NUM_UAV)
uav_speed = np.zeros(NUM_UAV)
uav_shots = np.zeros(NUM_UAV, dtype=int)
uav_last_release = np.full(NUM_UAV, -np.inf)


# 定义计划类
class ShotPlan:
    def __init__(self):
        self.uav = 0
        self.theta = 0
        self.v = 0
        self.t_rel = 0
        self.tau = 0
        self.t_det = 0
        self.D = np.zeros(3)
        self.C0 = np.zeros(3)


plan = []

# 预先计算"当前总目标值"（三枚导弹遮蔽时长之和）
mask1 = (dmin1 <= R)
mask2 = (dmin2 <= R)
mask3 = (dmin3 <= R)
curr_total = (np.sum(mask1) + np.sum(mask2) + np.sum(mask3)) * dt_eval

for chosen in range(MAX_TOTAL_SHOTS):
    best_gain = 0  # 本轮最大增益
    best_d1 = None
    best_d2 = None
    best_d3 = None
    best_shot = None  # 记录最佳候选

    for iu in range(NUM_UAV):
        if uav_shots[iu] >= 3:
            continue  # 每架最多 3 枚

        F0 = FY[iu, :]
        z0 = F0[2]

        # 候选航向角/速度：若该 UAV 尚未固定，可遍历网格；否则仅取固定值
        if not uav_fixed[iu]:
            ang_set = angles
            spd_set = speeds
        else:
            ang_set = [uav_theta[iu]]
            spd_set = [uav_speed[iu]]

        for ia in range(len(ang_set)):
            theta = ang_set[ia]
            uxF = np.cos(theta)
            uyF = np.sin(theta)

            for iv in range(len(spd_set)):
                vF = spd_set[iv]

                for id in range(len(t_grid)):
                    t_det = t_grid[id]

                    for it in range(len(tau_set)):
                        tau = tau_set[it]
                        t_rel = t_det - tau  # 投放时刻

                        # 约束：投放时刻>=0、与该 UAV 上一发间隔>=1s
                        if t_rel < 0:
                            continue
                        if t_rel < uav_last_release[iu] + 1.0:
                            continue

                        # 起爆点（水平=UAV(t_det)，竖直=自由落体）
                        x_det = F0[0] + vF * uxF * t_det
                        y_det = F0[1] + vF * uyF * t_det
                        z_det = z0 - 0.5 * g * tau ** 2

                        # 云团仅在 [t_det, t_det+20] 内有效
                        valid = (ts >= t_det) & (ts <= t_det + valid_win)
                        Cz = z_det - v_sink * (ts - t_det)

                        # === 计算该候选对三枚导弹的 d(t) ===
                        # M1
                        W1x = x_det - M1x
                        W1y = y_det - M1y
                        W1z = Cz - M1z

                        WU1 = W1x * U1x + W1y * U1y + W1z * U1z
                        lam1 = WU1 / UU1
                        lam1 = np.clip(lam1, 0, 1)

                        P1x = M1x + lam1 * U1x
                        P1y = M1y + lam1 * U1y
                        P1z = M1z + lam1 * U1z

                        d1 = np.sqrt((x_det - P1x) ** 2 + (y_det - P1y) ** 2 + (Cz - P1z) ** 2)
                        d1[~valid] = np.inf

                        # M2
                        W2x = x_det - M2x
                        W2y = y_det - M2y
                        W2z = Cz - M2z

                        WU2 = W2x * U2x + W2y * U2y + W2z * U2z
                        lam2 = WU2 / UU2
                        lam2 = np.clip(lam2, 0, 1)

                        P2x = M2x + lam2 * U2x
                        P2y = M2y + lam2 * U2y
                        P2z = M2z + lam2 * U2z

                        d2 = np.sqrt((x_det - P2x) ** 2 + (y_det - P2y) ** 2 + (Cz - P2z) ** 2)
                        d2[~valid] = np.inf

                        # M3
                        W3x = x_det - M3x
                        W3y = y_det - M3y
                        W3z = Cz - M3z

                        WU3 = W3x * U3x + W3y * U3y + W3z * U3z
                        lam3 = WU3 / UU3
                        lam3 = np.clip(lam3, 0, 1)

                        P3x = M3x + lam3 * U3x
                        P3y = M3y + lam3 * U3y
                        P3z = M3z + lam3 * U3z

                        d3 = np.sqrt((x_det - P3x) ** 2 + (y_det - P3y) ** 2 + (Cz - P3z) ** 2)
                        d3[~valid] = np.inf

                        # 与当前联合最小距离合并
                        nd1 = np.minimum(dmin1, d1)
                        nd2 = np.minimum(dmin2, d2)
                        nd3 = np.minimum(dmin3, d3)

                        # 目标：三导弹遮蔽时长之和
                        tot = (np.sum(nd1 <= R) + np.sum(nd2 <= R) + np.sum(nd3 <= R)) * dt_eval
                        gain = tot - curr_total

                        if gain > best_gain:
                            best_gain = gain
                            best_d1 = nd1.copy()
                            best_d2 = nd2.copy()
                            best_d3 = nd3.copy()

                            # 记录最佳候选
                            best_shot = ShotPlan()
                            best_shot.uav = iu
                            best_shot.theta = theta
                            best_shot.v = vF
                            best_shot.t_rel = t_rel
                            best_shot.tau = tau
                            best_shot.t_det = t_det
                            best_shot.D = np.array([
                                F0[0] + vF * uxF * t_rel,
                                F0[1] + vF * uyF * t_rel,
                                z0
                            ])
                            best_shot.C0 = np.array([x_det, y_det, z_det])

    # 若没有正增益，提前停止
    if best_gain <= 1e-3:
        print(f' 贪心迭代在第 {chosen + 1} 次停止：新增收益 {best_gain:.4f} s。')
        break

    # 接受本轮最优
    dmin1 = best_d1
    dmin2 = best_d2
    dmin3 = best_d3
    curr_total += best_gain

    iu = best_shot.uav
    if not uav_fixed[iu]:
        uav_fixed[iu] = True
        uav_theta[iu] = best_shot.theta
        uav_speed[iu] = best_shot.v

    uav_last_release[iu] = best_shot.t_rel
    uav_shots[iu] += 1
    plan.append(best_shot)

# ================= 结果统计（总时长/OR/AND） =================
mask1 = (dmin1 <= R)
mask2 = (dmin2 <= R)
mask3 = (dmin3 <= R)
sum_per_missile = np.array([np.sum(mask1), np.sum(mask2), np.sum(mask3)]) * dt_eval

OR_mask = mask1 | mask2 | mask3
AND_mask = mask1 & mask2 & mask3
OR_total = np.sum(OR_mask) * dt_eval
AND_total = np.sum(AND_mask) * dt_eval
TOTAL_sum = np.sum(sum_per_missile)

# ================= 打印策略 =================
print('========= 问题 5：5 架 UAV（每架≤3 枚）干扰 M1/M2/M3 的最优联合策略（贪心） =========')
# 先按 UAV 编号分组输出
for iu in range(NUM_UAV):
    idx = [i for i, p in enumerate(plan) if p.uav == iu]
    if not idx:
        continue

    print(f'—— FY{iu + 1} —— 固定航向θ={np.rad2deg(uav_theta[iu]):.2f}°，固定速度 v={uav_speed[iu]:.2f} m/s')
    for k, shot_idx in enumerate(idx):
        s = plan[shot_idx]
        print(f' 第{k + 1:2d} 枚：投放 t_rel={s.t_rel:.3f} s，τ={s.tau:.3f} s，起爆 t_det={s.t_det:.3f} s')
        print(f'    D=({s.D[0]:.3f}, {s.D[1]:.3f}, {s.D[2]:.3f})  C0=({s.C0[0]:.3f}, {s.C0[1]:.3f}, {s.C0[2]:.3f})')

print(
    f'——遮蔽总时长（分导弹）: M1={sum_per_missile[0]:.3f} s, M2={sum_per_missile[1]:.3f} s, M3={sum_per_missile[2]:.3f} s')
print(f'—— OR 覆盖(任一被遮蔽): {OR_total:.3f} s | AND 覆盖(同时遮蔽三枚): {AND_total:.3f} s')
print(f'——目标函数值（和计） : TOTAL = {TOTAL_sum:.3f} s')